import cv2
import numpy as np
import matplotlib.pyplot as plt


def edge_detection(img, method='sobel'):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img.copy(), cv2.COLOR_RGB2GRAY)
    height, width = img.shape  # 确定height和weight以确认循环次数
    if method == 'sobel':
        # sobel算子
        sx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
        sy = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
    elif method == 'prewitt':
        # 更换算子为prewitt
        sx = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])
        sy = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
    else:
        print("不存在 %s 滤波核" % method)
        return
    # 初始化矩阵
    dSobel = np.zeros((height, width))
    dSobelx = np.zeros((height, width))
    dSobely = np.zeros((height, width))
    # Gx，Gy分别为x,y方向的求偏导，通过卷积实现（数字近似）
    Gx = np.zeros(img.shape)
    Gy = np.zeros(img.shape)
    # 通过两层循环实现卷积操作
    for i in range(height - 2):
        for j in range(width - 2):
            Gx[i + 1, j + 1] = abs(np.sum(img[i:i + 3, j:j + 3] * sx))  # 求当前位置的gx
            Gy[i + 1, j + 1] = abs(np.sum(img[i:i + 3, j:j + 3] * sy))  # 求当前位置的gy
            dSobel[i + 1, j + 1] = abs(Gx[i + 1, j + 1]) + abs(Gy[i + 1, j + 1])  # 合成梯度近似
            dSobelx[i + 1, j + 1] = np.sqrt(Gx[i + 1, j + 1])
            dSobely[i + 1, j + 1] = np.sqrt(Gy[i + 1, j + 1])
    # 截断为Uint8
    a = np.clip(dSobelx, 0, 255).astype(np.uint8)
    b = np.clip(dSobely, 0, 255).astype(np.uint8)
    c = np.clip(dSobel, 0, 255).astype(np.uint8)
    return a, b, c


if __name__ == '__main__':
    img = cv2.imread('4.jpg')
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 将图像转化为灰度图
    a, b, c = edge_detection(gray)
    plt.subplot(321), plt.imshow(img), plt.title('original_gray')
    plt.subplot(322), plt.imshow(a, cmap=plt.cm.gray), plt.title('x')
    plt.subplot(323), plt.imshow(b, cmap=plt.cm.gray), plt.title('y')
    plt.subplot(324), plt.imshow(c, cmap=plt.cm.gray), plt.title('both')
    plt.show()
